# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os
import math
import datetime

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from coin.base.datetime_util import to_datetime
from coin.support.feed_tool.feed_stats.logic.util import datetime_span
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.database.base.h5py_io import h5py_load

FLAGS = flags.FLAGS


def plot_df(df, csv_root, trading_date_str, full_symbol):
  fig, ax1 = plt.subplots(figsize=(20, 10))

  plt.title('%s %s open interest VS mid_price' % (trading_date_str, full_symbol))
  ax1.plot(df.index, df[full_symbol], 'r')
  ax1.set_ylabel('open_interest', color='r')
  tick_space = (ax1.get_yticks()[-1] - ax1.get_yticks()[0]) / 20
  tick_space = int(math.ceil(tick_space / 5)) * 5
  ax1.yaxis.set_major_locator(ticker.MultipleLocator(tick_space))
  ax1.grid(True)

  ax2 = ax1.twinx()
  ax2.plot(df.index, df['mid_price'], color='g', marker='.')
  ax2.set_ylabel('mid_price', color='g')
  plt.savefig(
      os.path.join(csv_root,
                   '%s_%s_open_interest_vs_mid_price.png' % (trading_date_str, full_symbol)))
  plt.close()


def sample_df_per_second(df):
  return df.resample('S', closed='left', label='left').ffill()


def read_csv_into_df(path):
  df = pandas.read_csv(path, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['timestamp'], unit='ns')
  df = df.set_index('datetime')
  assert len(df) > 0, 'Empty instrument csv: %s' % path
  return df


def read_hdf5_into_df(path):
  m = h5py_load(path)

  df = pandas.DataFrame(data=m)
  return df


def _gen_products_symbols(base, trading_date):
  products = []
  full_symbols = []

  product = BitmexFuturesProduct.FromStr('%s-USD.PERPETUAL' % base)
  products.append(product)
  full_symbols.append(product.full_symbol)

  product = OkexFuturesProduct.FromStr('%s-USD.QUARTER' % base, trading_date)
  products.append(product)
  full_symbols.append(product.full_symbol)

  price_product = product
  products.append(product)
  full_symbols.append(product.full_symbol_v3)

  product = OkexFuturesProduct.FromStr('%s-USD-SWAP' % base)
  products.append(product)
  full_symbols.append(product.full_symbol_v3)

  return products, full_symbols, price_product


def main(argv):
  hdf5_root = FLAGS.hdf5_root
  assert hdf5_root, '--hdf5_root must be specified.'
  csv_root = FLAGS.csv_root
  assert csv_root, '--csv_root must be specified.'

  trading_date_str = FLAGS.trading_date
  trading_date = datetime.datetime.strptime(trading_date_str, '%Y%m%d')
  base = flags.FLAGS.base

  products, full_symbols, price_product = _gen_products_symbols(base, trading_date)

  csv_book = 'book/%s.csv' % price_product.full_symbol
  hdf5_open_interest_open = \
      'futures/open_interest/Open_open_interest.value--align-ti--interval-1min--type-matrix.hdf5'
  hdf5_open_interest_interval = \
      'futures/open_interest/intervals.value--align-i--interval-1min--type-vector.hdf5'
  hdf5_full_symbol = \
      'presto/metadata/symbols.value--align-i--type-vector.hdf5'

  book_df = read_csv_into_df(os.path.join(csv_root, csv_book))
  book_df = book_df.resample('T', closed='left', label='left').ffill()
  symbol_columns_df = read_hdf5_into_df(os.path.join(hdf5_root, trading_date_str, hdf5_full_symbol))
  symbol_columns = symbol_columns_df.iloc[:, 0].apply(lambda x: x.decode()).tolist()
  open_interest_open_df = read_hdf5_into_df(
      os.path.join(hdf5_root, trading_date_str, hdf5_open_interest_open))
  open_interest_open_df.columns = symbol_columns
  interval_df = read_hdf5_into_df(
      os.path.join(hdf5_root, trading_date_str, hdf5_open_interest_interval))
  timestamp_column = interval_df.iloc[:, 0]
  open_interest_open_df['timestamp'] = timestamp_column
  open_interest_open_df['datetime'] = pandas.to_datetime(open_interest_open_df['timestamp'],
                                                         unit='ns')
  open_interest_open_df = open_interest_open_df.set_index('datetime')
  df = pandas.merge(book_df, open_interest_open_df, how='inner', left_index=True, right_index=True)
  df['mid_price'] = (df['bid0'] + df['ask0']) / 2

  # Put output plot into same csv_root directory.
  for full_symbol in full_symbols:
    plot_df(df, csv_root, trading_date_str, full_symbol)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('hdf5_root', None, 'Input hdf5 root directory.')

  flags.DEFINE_string('csv_root', None, 'Input csv root directory.')

  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('base', 'BTC', 'symbol base name')

  app.run(main)
